#include "acl/acl.h"
#include "ascend_utils.h"

namespace knowhere {
namespace ascend {
int
ascend_device_manager::random_choose() const {
    srand(time(NULL));
    int device_id = rand() % memory_load_.size();
    LOG_KNOWHERE_INFO_ << "random_choose device: " << device_id;
    return device_id;
}

int
ascend_device_manager::choose_with_load(size_t load) {
    std::lock_guard<std::mutex> lock(mtx_);

    auto it = std::min_element(memory_load_.begin(), memory_load_.end());
    *it += load;
    int device_id = std::distance(memory_load_.begin(), it);
    LOG_KNOWHERE_INFO_ << "min choose_with_load device: " << device_id << ", load:" << load;
    return device_id;
}

void
ascend_device_manager::release_load(int device_id, size_t load) {
    if (size_t(device_id) < memory_load_.size()) {
        std::lock_guard<std::mutex> lock(mtx_);
        memory_load_[device_id] -= load;
        LOG_KNOWHERE_INFO_ << "release_load device: " << device_id << ", load:" << load;
    } else {
        LOG_KNOWHERE_WARNING_ << "please check device id " << device_id;
    }
}

ascend_device_manager::ascend_device_manager() {}

void ascend_device_manager::init() {
    std::lock_guard<std::mutex> lock(mtx_);
    if (inited_.load()) {
        return;
    }

    uint32_t device_counts;
    auto ret = aclrtGetDeviceCount(&device_counts);
    if (ret != ACL_SUCCESS) {
        LOG_KNOWHERE_FATAL_ << "failed to get Ascend Device count!";
        KNOWHERE_ASCEND_THROW_MSG("failed to get Ascend Device count!");
    }
    memory_load_.resize(device_counts);
    std::fill(memory_load_.begin(), memory_load_.end(), 0);

    device_mutex_ptres.resize(device_counts);
    for (uint32_t dev = 0; dev < device_counts; dev++) {
        device_mutex_ptres[dev] = new std::mutex();
    }

    inited_.store(true);
}

std::mutex* ascend_device_manager::get_dev_mutex_ptr(int32_t device_id) {
    if (device_id < 0 || static_cast<size_t>(device_id + 1) > device_mutex_ptres.size()) {
        LOG_KNOWHERE_ERROR_ << "bad device_id:" << device_id;
        KNOWHERE_ASCEND_THROW_FMT("device_id:%i is out of available devices!", device_id);
    }

    return device_mutex_ptres[device_id];
}

ascend_device_manager&
ascend_device_manager::instance() {
    static ascend_device_manager mgr;
    if (!mgr.inited_.load()) {
        mgr.init();
    }
    return mgr;
}

} // namespace ascend
}